/*
 * Copyright (C) Hisilicon Technologies Co., Ltd. 2018-2019. All rights reserved.
 * Description:
 * Author:
 * Create: 2018-05-19
 */

#include <stdio.h>
#include <stdlib.h>
#include "string.h"
#ifdef ON_BOARD
#include "mpi_sys.h"
#include "mpi_vb.h"
#else
#include "hi_comm_svp.h"
#include "hi_nnie.h"
#include "mpi_nnie.h"
#endif
#include "math.h"
#include "sample_log.h"
#include "sample_save_blob.h"
#include "sample_data_utils.h"
#include "sample_runtime_classify.h"
#include "sample_runtime_detection_rfcn.h"
#include "sample_runtime_detection_ssd.h"

static HI_VOID SAMPLE_RUNTIME_Usage(const char *pchPrgName)
{
    SAMPLE_LOG_INFO("Usage : %s <index>", pchPrgName);
    SAMPLE_LOG_INFO("index:");
    SAMPLE_LOG_INFO("\t 0) Alexnet(VNode)");
    SAMPLE_LOG_INFO("\t 1) RFCN(VNode->RNode->VNode)");
    SAMPLE_LOG_INFO("\t 2) SSD)");
}

#ifdef ON_BOARD
static HI_S32 SAMPLE_COMM_SVP_SysInit(HI_VOID)
{
    HI_MPI_SYS_Exit();
    HI_MPI_VB_Exit();

    VB_CONFIG_S struVbConf = {0};
    struVbConf.u32MaxPoolCnt = 2;
    struVbConf.astCommPool[1].u64BlkSize = 768 * 576 * 2;
    struVbConf.astCommPool[1].u32BlkCnt = 1;

    HI_S32 s32Ret = HI_MPI_VB_SetConfig((const VB_CONFIG_S *)&struVbConf);

    if (s32Ret != HI_SUCCESS) {
        SAMPLE_LOG_INFO("HI_MPI_VB_SetConfig error");
        return HI_FAILURE;
    }

    s32Ret = HI_MPI_VB_Init();

    if (s32Ret != HI_SUCCESS) {
        SAMPLE_LOG_INFO("HI_MPI_VB_Init error");
        return HI_FAILURE;
    }

    s32Ret = HI_MPI_SYS_Init();

    if (s32Ret != HI_SUCCESS) {
        SAMPLE_LOG_INFO("HI_MPI_SYS_Init error");
        return HI_FAILURE;
    }

    return s32Ret;
}

static HI_S32 SAMPLE_COMM_SVP_SysExit(HI_VOID)
{
    HI_S32 s32Ret = HI_MPI_SYS_Exit();

    if (s32Ret != HI_SUCCESS) {
        SAMPLE_LOG_INFO("HI_MPI_SYS_Exit error");
        return HI_FAILURE;
    }

    s32Ret = HI_MPI_VB_Exit();

    if (s32Ret != HI_SUCCESS) {
        SAMPLE_LOG_INFO("HI_MPI_VB_Exit error");
        return HI_FAILURE;
    }

    return HI_SUCCESS;
}
#endif

HI_S32 HandleSampleNet(HI_CHAR* binName, HI_S32 index)
{
    HI_S32 ret = HI_FAILURE;

    switch (index) {
        case '0':
            ret = SAMPLE_AlexNet();
            break;
        case '1':
            ret = SAMPLE_RFCN();
            break;
        case '2':
            ret = SAMPLE_SSD();
            break;
        default:
            SAMPLE_RUNTIME_Usage(binName);
            break;
    }

    return ret;
}

HI_S32 main(HI_S32 argc, HI_CHAR* argv[])
{
    if ((argc != 2) || (*argv[1] < '0') || (*argv[1] > '2') || strlen(argv[1]) != 1) {
        SAMPLE_RUNTIME_Usage(argv[0]);
        return HI_FAILURE;
    }

    if (!strncmp(argv[1], "-h", 2)) {
        SAMPLE_RUNTIME_Usage(argv[0]);
        return HI_SUCCESS;
    }

    HI_CHAR* selcetNet[] = { "alexnet", "rfcn", "ssd" };
    if (atoi(argv[1]) < sizeof(selcetNet) / sizeof(selcetNet[0])) {
        SAMPLE_LOG_INFO("\n\n====================start to run %s====================\n", selcetNet[atoi(argv[1])]);
    }

#ifdef ON_BOARD
    HI_S32 s32Ret = SAMPLE_COMM_SVP_SysInit();
    SAMPLE_CHK_RET(s32Ret != HI_SUCCESS, HI_FAILURE, "SAMPLE_COMM_SVP_SysInit fail");
    s32Ret = HandleSampleNet(argv[0], *argv[1]);
#else
    HI_S32 s32Ret = HandleSampleNet(argv[0], *argv[1]);
#endif

    SAMPLE_CHK_PRINTF((s32Ret == HI_SUCCESS), "\n==========run %s success==========\n", selcetNet[atoi(argv[1])]);
    SAMPLE_CHK_PRINTF((s32Ret != HI_SUCCESS), "\n==========run %s fail==========\n", selcetNet[atoi(argv[1])]);

#ifdef ON_BOARD
    s32Ret = SAMPLE_COMM_SVP_SysExit();
    SAMPLE_CHK_RET(s32Ret != HI_SUCCESS, HI_FAILURE, "SAMPLE_COMM_SVP_SysExit fail");
#endif

    return s32Ret;
}
